/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include <iostream>

#include "acl/acl.h"
#include "atlas_model.h"
#include "do_process.h"
#include <sys/time.h>
#include "thread_safe_queue.h"
using namespace std;

 // 模型固定输入大小和模型路径
namespace {
    const char * kAclConfigPath = "../src/acl.json";
    DoProcess::ModelInfoMap infoMap{
        // algorithm name -> { model path, width, height }
        {
            "sort", 
            {
                { "../model/yolo/yolov5s_trained.om", 640, 640 } 
            }
        },
        {
            "deepsort",
            {
                { "../model/yolo/yolov5s_trained.om", 640, 640 },
                { "../model/deepsort/deep.om", 64, 128 }
            }
        },
        {
            "retina",
            {
                {"../model/retina/RetinaNet_sim_608.om", 608, 608 }
            }
        },
        {
            "dimp",
            {
                {"../model/dimp/dimp-InitialModel.om", 0, 0},
                {"../model/dimp/dimp-ScoreModel.om", 0, 0},
                {"../model/dimp/dimp-IoUModel.om", 0, 0},
                {"../model/dimp/dimp-Optimizer.om", 0, 0},
                {"../model/dimp/dimp-UpdateModel.om", 0, 0},
                {"../data/dimp/test6/groundtruth_rect.txt", 0, 0}
            }
        }
    };
}

DoProcess::DoProcess()
:deviceId_(0), context_(nullptr), stream_(nullptr), isInited_(false)
{
    initMap_.insert(std::make_pair("deepsort", std::bind(&DoProcess::InitDeepSORT, this, std::placeholders::_1)));
    initMap_.insert(std::make_pair("sort", std::bind(&DoProcess::InitSORT, this, std::placeholders::_1)));
    initMap_.insert(std::make_pair("dimp", std::bind(&DoProcess::InitDiMP, this, std::placeholders::_1)));
    initMap_.insert(std::make_pair("retina", std::bind(&DoProcess::InitRetina, this, std::placeholders::_1)));

    processMap_.insert(std::make_pair("deepsort", std::bind(&DoProcess::ProcessDeepSORT, this, std::placeholders::_1)));
    processMap_.insert(std::make_pair("sort", std::bind(&DoProcess::ProcessSORT, this, std::placeholders::_1)));
    processMap_.insert(std::make_pair("dimp", std::bind(&DoProcess::ProcessDiMP, this, std::placeholders::_1)));
    processMap_.insert(std::make_pair("retina", std::bind(&DoProcess::ProcessRetina, this, std::placeholders::_1)));

    outputSet_.insert("text");
    outputSet_.insert("video");
    outputSet_.insert("image");
    outputSet_.insert("network_image");
    outputSet_.insert("network_result");
    outputSet_.insert("");
}

DoProcess::~DoProcess() {
    DestroyResource();
}

void DoProcess::DestroyResource()
{
    if (yolo_ != nullptr) yolo_->DestroyResource();
    if (deepsort_ != nullptr) deepsort_->DestroyResource();
    if (sort_ != nullptr) sort_->DestroyResource();
    
    storage_->DestroyResource();
    sender_->DestroyResource();

    aclError ret;
    if (stream_ != nullptr) {
        ret = aclrtDestroyStream(stream_);
        if (ret != ACL_ERROR_NONE) {
            ATLAS_LOG_ERROR("destroy stream failed");
        }
        stream_ = nullptr;
    }
    INFO_LOG("end to destroy stream");

    if (context_ != nullptr) {
        ret = aclrtDestroyContext(context_);
        if (ret != ACL_ERROR_NONE) {
            ERROR_LOG("destroy context failed");
        }
        context_ = nullptr;
    }
    INFO_LOG("end to destroy context");

    ret = aclrtResetDevice(deviceId_);
    if (ret != ACL_ERROR_NONE) {
        ATLAS_LOG_ERROR("reset device failed\n");
    }
    ATLAS_LOG_INFO("end to reset device is %d\n", deviceId_);

    ret = aclFinalize();
    if (ret != ACL_ERROR_NONE) {
        ATLAS_LOG_ERROR("finalize acl failed\n");
    }
    ATLAS_LOG_INFO("end to finalize acl");

}

AtlasError DoProcess::InitResource() {
    // ACL init
    aclError ret = aclInit(kAclConfigPath);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl init failed");
        return FAILED;
    }
    INFO_LOG("acl init success");

    // open device
    ret = aclrtSetDevice(deviceId_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl open device %d failed", deviceId_);
        return FAILED;
    }
    INFO_LOG("open device %d success", deviceId_);

    // create context (set current)
    ret = aclrtCreateContext(&context_, deviceId_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl create context failed");
        return FAILED;
    }
    INFO_LOG("create context success");

    // create stream
    ret = aclrtCreateStream(&stream_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl create stream failed");
        return FAILED;
    }
    INFO_LOG("create stream success");

    ret = aclrtGetRunMode(&runMode_);
    if (ret != ACL_ERROR_NONE) {
        ERROR_LOG("acl get run mode failed");
        return FAILED;
    }
    return ATLAS_OK;
}

AtlasError DoProcess::InitStorage(uint32_t width, uint32_t height) {
    
    storage_ = make_shared<TargetStorage>();
    auto dest = params_.destination;
    // 如果当前输出format是网络模式就不采纳params_.destination作为输出路径，因为此时该值表示的是网络地址
    if (params_.format.find("network") != params_.format.npos) {
        dest = "";
    }
    storage_->init(width, height, dest, context_);
    return ATLAS_OK;
}

AtlasError DoProcess::Init(const DoProcess::processParams& params) {
    if (isInited_) {
        ATLAS_LOG_INFO("instance is initied already!\n");
        return ATLAS_OK;
    }

    params_ = params;

    // 检查format的合法性
    if (outputSet_.count(params.format) == 0) {
        ATLAS_LOG_INFO("Don't support %s output format!", params_.format.c_str());
        return FAILED;
    }

    // 检查task的合法性
    if (initMap_.count(params_.task) == 0) {
        ATLAS_LOG_INFO("Don't support task %s!", params_.task.c_str());
        return FAILED;
    }

    // 初始化系统资源
    AtlasError ret = InitResource();
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init acl resource failed, error: %d", ret);
        return ret;
    }

    // 初始化输入任务
    ret = initMap_[params_.task](std::ref(infoMap[params_.task]));
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init task %s failed, error: %d", params_.task.c_str(), ret);
        return ret;
    }

    // 初始化网络发送器
    if (params_.format.find("network") != params_.format.npos) {
        sender_ = make_shared<TargetSender>();
        ret = (AtlasError)sender_->Init(context_, stream_, params_.destination);
        if (ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("Init network sender failed\n");
            return ret;
        }
    }
    
    isInited_ = true;
    return ATLAS_OK;
}

AtlasError DoProcess::Process(ImageData& image) {

    // 第一次处理时进行一次持久化器的初始化（必须在此进行，因为持久化器依赖输入图像的宽高）
    std::call_once(onceFlag_, &DoProcess::InitStorage, this, image.width, image.height);

    // 执行任务
    return processMap_[params_.task](std::ref(image));
}

AtlasError DoProcess::InitSORT(ModelInfos& infos) {

    if (infos.size() < 1) {
        ATLAS_LOG_ERROR("InitSORT: Insufficient model informations\n");
        return ATLAS_ERROR;
    }

    yolo_ = make_shared<YoloObjectDetect>(infos[0].modelPath.c_str(), infos[0].width, infos[0].height);
    sort_ = make_shared<SORT>();

    AtlasError ret = yolo_->Init(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init yolo dectector failed\n");
        return ret;
    }

    ret = sort_->Init(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init sort mot tracker failed\n");
        return ret;
    }
    return ATLAS_OK;
}

AtlasError DoProcess::InitDeepSORT(ModelInfos& infos) {

    if (infos.size() < 2) { 
        ATLAS_LOG_ERROR("InitDeepSORTSORT: Insufficient model informations\n");
        return ATLAS_ERROR;
    }

    yolo_ = make_shared<YoloObjectDetect>(infos[0].modelPath.c_str(), infos[0].width, infos[0].height);
    deepsort_ = make_shared<DeepSort>(infos[1].modelPath.c_str(), infos[1].width, infos[1].height);

    AtlasError ret = yolo_->Init(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init yolo dectector failed\n");
        return ret;
    }

    ret = deepsort_->Init(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init deepsort_ mot tracker failed\n");
        return ret;
    }
    return ATLAS_OK;
}

AtlasError DoProcess::InitDiMP(ModelInfos& infos) {

    if (infos.size() < 6) {
        ATLAS_LOG_ERROR("InitSORT: Insufficient model informations\n");
        return ATLAS_ERROR;
    }
    const char* dimp_modelpath[] = {infos[0].modelPath.c_str(), infos[1].modelPath.c_str(), infos[2].modelPath.c_str(), infos[3].modelPath.c_str(), infos[4].modelPath.c_str()};

    dimp_ = make_shared<Tracker>(dimp_modelpath, infos[5].modelPath.c_str());

    AtlasError ret = dimp_->LoadAllModel(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("Init deepsort_ mot tracker failed\n");
        return ret;
    }

    return ATLAS_OK;
}

AtlasError DoProcess::InitRetina(ModelInfos& infos) {

    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(1,&cpuset);
    if (pthread_setaffinity_np(pthread_self(),sizeof(cpu_set_t), &cpuset) != 0) {
        ATLAS_LOG_ERROR("InitRetina: pthread_setaffinity_np error\n");
        return ATLAS_ERROR;
    }

    if (infos.size() < 1) {
        ATLAS_LOG_ERROR("InitRetina: Insufficient model informations\n");
        return ATLAS_ERROR;
    }

    retina_ = make_shared<RetinaObjectDetect>(infos[0].modelPath.c_str(), infos[0].width, infos[0].height);

    AtlasError ret = retina_->Init(deviceId_, context_, stream_, runMode_);
    if (ret != ATLAS_OK) {
        ATLAS_LOG_ERROR("InitRetina: Init dectector failed\n");
        return ret;
    }
    return ATLAS_OK;
}

AtlasError DoProcess::ProcessSORT(ImageData& image) {

    static bool isSetSORT = false;
    if (!isSetSORT) {
        sort_->Set(image.width, image.height);
        isSetSORT = true;
    }

    cv::Mat bgr;
    vector<BBox> dectections;
    vector<BBox> trackResults;
    AtlasError ret;

    {
        time_guard t("Detection");
        ret = yolo_->Process(image, bgr, dectections);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("object detection failed\n");
            return ATLAS_ERROR;
        }
    }
    

    {
        time_guard t("Track");
        ret = sort_->Process(bgr, dectections, trackResults);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("deepsort failed\n");
            return ATLAS_ERROR;
        }

    }
        
    {

        time_guard t("Output");
        // --------------同步操作----------------------
        // storage_->update(trackResults, image.id, bgr);   // 记录轨迹和snapshot
        // storage_->draw_detections(bgr, dectections);    // 画出黑色的检测框
        // storage_->draw_trajectories(bgr, trackResults, false);   // 画出彩色的跟踪框
        // storage_->save_as_txt(dectections, image.name, image.width, image.height, true);  // 结果存为txt
        // storage_->dvpp_write_video(bgr);                        // 用dvpp编码为视频
        // storage_->dvpp_write_image(bgr, image.id);              // 用dvpp保存为图片
        // storage_->cv_write_video(bgr);                       // 用opencv编码为视频
        // storage_->cv_write_image(bgr, image.id);             // 用opencv保存为图片
        // --------------多线程操作----------------------
        // storage_->StorageDataEnQueue(bgr, image.id);             // 编码
        // sender_->SendDataEnQueue(trackResults, image.id, bgr);   // 发送图片到上位机
        // sender_->SendDataEnQueue(trackResults, image.id);        // 只发送检测结果到上位机

        sort_->DrawTrajectories(bgr, trackResults, false);   // 画出彩色的跟踪框
        // yolo_->DrawDetections(bgr, dectections);
        // cout << dectections.size() << endl;

        if (params_.format == "network_image") {
            sender_->SendDataEnQueue(trackResults, image.id, bgr);   // 发送图片到上位机
        } else if (params_.format == "network_result") {
            sender_->SendDataEnQueue(trackResults, image.id);        // 只发送检测结果到上位机
        } else if (params_.format == "text") {
            storage_->save_as_txt(dectections, image.name, image.width, image.height, true);  // 结果存为txt(yolo形式)
        } else if (params_.format == "video") {
            storage_->StorageDataEnQueue(bgr, image.id);             // 编码为视频
        } else if (params_.format == "image") {
            storage_->dvpp_write_image(bgr, image.id);              // 用dvpp保存为图片
        }
    }

    INFO_LOG("Frame id: %d", image.id);
    return ATLAS_OK;
}

AtlasError DoProcess::ProcessDeepSORT(ImageData& image) {

    static bool isSetDeepSORT = false;
    if (!isSetDeepSORT) {
        deepsort_->Set(image.width, image.height);
        isSetDeepSORT = true;
    }

    AtlasError ret;
    cv::Mat bgr;
    vector<BBox> dectections;
    vector<BBox> trackResults;

    {
        time_guard t("Detection");
        ret = yolo_->Process(image, bgr, dectections);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("object detection failed\n");
            return ATLAS_ERROR;
        }
    }
        
    {
        time_guard t("Track");
        ret = deepsort_->Process(bgr, dectections, trackResults);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("deepsort failed\n");
            return ATLAS_ERROR;
        }
    }

    {
        time_guard t("Output");

        deepsort_->DrawTrajectories(bgr, trackResults, false);   // 画出彩色的跟踪框

        if (params_.format == "network_image") {
            sender_->SendDataEnQueue(trackResults, image.id, bgr);   // 发送图片到上位机
        } else if (params_.format == "network_result") {
            sender_->SendDataEnQueue(trackResults, image.id);        // 只发送检测结果到上位机
        } else if (params_.format == "text") {
            storage_->save_as_txt(dectections, image.name, image.width, image.height, true);  // 结果存为txt(yolo形式)
        } else if (params_.format == "video") {
            storage_->StorageDataEnQueue(bgr, image.id);             // 编码为视频
        } else if (params_.format == "image") {
            storage_->dvpp_write_image(bgr, image.id);              // 用dvpp保存为图片
        }
    }


    INFO_LOG("Frame id: %d", image.id);
    return ATLAS_OK;
}

AtlasError DoProcess::ProcessDiMP(ImageData& image) {
    AtlasError ret;
    cv::Mat img;
    cv::Mat yuvimg(image.height * 3 / 2, image.width, CV_8UC1, image.data.get());
    cv::cvtColor(yuvimg, img, cv::COLOR_YUV2BGR_NV12);
    Target_Box result;
    vector<BBox> DimpRect;
    BBox b;

    if (image.id == 0) {
        ret = dimp_->InitFirstFrame(img);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("InitFirstFrame failed\n");
            return ATLAS_ERROR;
        }
    }

    else {
        ret = dimp_->Tracking(img, result);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("Tracking failed\n");
            return ATLAS_ERROR;
        }
        b.rect.x = result.pos.x;
        b.rect.y = result.pos.y;
        b.rect.width = result.size.x;
        b.rect.height = result.size.y;
        DimpRect.push_back(b);

        Utils::draw_bbox(img, b.rect);
        if (params_.format == "network_image") {
            sender_->SendDataEnQueue(DimpRect, image.id, img);   // 发送图片到上位机
        } else if (params_.format == "network_result") {
            sender_->SendDataEnQueue(DimpRect, image.id);        // 只发送检测结果到上位机
        } else if (params_.format == "text") {
            storage_->save_as_txt(DimpRect, image.name, image.width, image.height, true);  // 结果存为txt(yolo形式)
        } else if (params_.format == "video") {
            storage_->StorageDataEnQueue(img, image.id);             // 编码为视频
        } else if (params_.format == "image") {
            storage_->dvpp_write_image(img, image.id);              // 用dvpp保存为图片
        }
    }

    return ATLAS_OK;
}

AtlasError DoProcess::ProcessRetina(ImageData& image) {

    AtlasError ret;
    cv::Mat bgr;
    vector<BBox> dectections;

    {
        time_guard t("Detection");
        ret = retina_->Process(image, bgr, dectections);
        if(ret != ATLAS_OK) {
            ATLAS_LOG_ERROR("object detection failed\n");
            return ATLAS_ERROR;
        }
    }

    {
        time_guard t("Output");

        retina_->DrawBoundBoxToImage(dectections, bgr);

        if (params_.format == "network_image") {
            sender_->SendDataEnQueue(dectections, image.id, bgr);   // 发送图片到上位机
        } else if (params_.format == "network_result") {
            sender_->SendDataEnQueue(dectections, image.id);        // 只发送检测结果到上位机
        } else if (params_.format == "text") {
            storage_->save_as_txt(dectections, image.name, image.width, image.height, true);  // 结果存为txt(yolo形式)
        } else if (params_.format == "video") {
            storage_->StorageDataEnQueue(bgr, image.id);             // 编码为视频
        } else if (params_.format == "image") {
            storage_->dvpp_write_image(bgr, image.id);              // 用dvpp保存为图片
        }
    }

    
    return ATLAS_OK;
}

